import os
import csv
from concurrent.futures import ThreadPoolExecutor, Future
from typing import List, Tuple, Set, Union, Dict
from models import Sample
from event_bus import EventBus
from PIL import Image
from utils import switch

pool = ThreadPoolExecutor(max_workers=os.cpu_count() * 5 if os.cpu_count() is not None else 5)
cates = Sample.categories.children().alias("cates")


class BaseGrouping:
    export_location: str
    category_file: str
    sample_file: str
    validation_file: str
    resize_to: int
    padding: bool
    sample_image_type: str = "png"
    categories: Set[str]
    split_groups: int
    training_ratio: float
    validate_ratio: float
    groups: List[Union[Tuple[Set[str], Set[str]], Dict[str, Tuple[Set[str], Set[str]]]]] = list()
    task_record: List[Future] = list()

    def __init__(self, export_location: str, category_file: str = "category", sample_file: str = "sample",
                 validation_file: str = "validation", resize_to: int = 0, padding: bool = False,
                 sample_type: str = "png", categories: Set[str] = None, group: int = 1, training: float = 0.,
                 validate: float = 0.):
        self.export_location = export_location
        self.category_file = category_file
        self.sample_file = sample_file
        self.validation_file = validation_file
        self.resize_to = resize_to
        self.padding = padding
        self.sample_image_type = sample_type.lower()
        self.categories = list(categories) if categories is not None else list()
        self.split_groups = group
        self.training_ratio = training
        self.validate_ratio = validate

    def write_categories(self):
        with open(os.path.join(self.export_location, f"{self.category_file}.csv"), 'w') as file:
            writer = csv.writer(file)
            writer.writerow(list(self.categories))

    def calculate_works(self) -> int:
        pass

    def split_group(self):
        pass

    def write_group_description(self, group: int):
        pass

    def generate_samples(self):
        pass

    def save_sample_picture(self, target_location: str, sample: Sample):
        try:
            source_image = Image.open(sample.source)
            cropped_image = source_image.crop((max([0, sample.crop_x]),
                                               max([0, sample.crop_y]),
                                               min([sample.crop_x + sample.crop_width, source_image.width]),
                                               min([sample.crop_y + sample.crop_height, source_image.height])))
            if self.padding:
                longer_side = max([sample.crop_width, sample.crop_height])
                new_image = Image.new(mode="RGBA", size=(longer_side, longer_side))
                dest_x = (new_image.width - cropped_image.width) // 2
                dest_y = (new_image.height - cropped_image.height) // 2
            else:
                new_image = Image.new(mode="RGBA", size=(sample.crop_width, sample.crop_height))
                dest_x, dest_y = 0, 0
            new_image.alpha_composite(cropped_image.convert(mode="RGBA"), (dest_x, dest_y))
            if self.resize_to > 0:
                ratio = self.resize_to / max([new_image.width, new_image.height])
                resized = new_image.resize((round(new_image.width * ratio), round(new_image.height * ratio)),
                                           Image.LANCZOS)
            else:
                resized = new_image
            for case in switch(self.sample_image_type):
                if case("png"):
                    resized.save(os.path.join(target_location, f"{sample.target_fingerprint}.png"), "PNG")
                if case("jpg") or case("jpeg"):
                    resized.convert(mode="RGB") \
                        .save(os.path.join(target_location, f"{sample.target_fingerprint}.jpg"), "JPEG")
        except (OSError, ValueError, IOError) as e:
            EventBus().post_topic(event=e, topic="export_failed")
            return
        EventBus().post_topic(event=None, topic="export_job_done")

    def cancel(self):
        for task in self.task_record:
            task.cancel()
